package com.idata.common.Utils;

import com.kennycason.kumo.CollisionMode;
import com.kennycason.kumo.WordCloud;
import com.kennycason.kumo.WordFrequency;
import com.kennycason.kumo.bg.CircleBackground;
import com.kennycason.kumo.font.KumoFont;
import com.kennycason.kumo.font.scale.SqrtFontScalar;
import com.kennycason.kumo.nlp.FrequencyAnalyzer;
import com.kennycason.kumo.nlp.tokenizers.ChineseWordTokenizer;
import com.kennycason.kumo.palette.ColorPalette;

import java.awt.*;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.ThreadLocalRandom;
import java.util.stream.Collectors;

/**
 * @File WordCloudUtils.Java
 * @Modify Time      @Author    @Version    @Description
 * ------------      -------    --------    -----------
 * 2023/9/22         yd.j      1.0         None
 */
public class WordCloudUtils {

    public static OutputStream createWordCloud(List<String> data) throws IOException {
        final FrequencyAnalyzer frequencyAnalyzer = new FrequencyAnalyzer();
        frequencyAnalyzer.setWordFrequenciesToReturn(600);
        frequencyAnalyzer.setMinWordLength(2);
        frequencyAnalyzer.setWordTokenizer(new ChineseWordTokenizer());
        // 可以直接从文件中读取
        final List<WordFrequency> wordFrequencies = new ArrayList<>();

        // 词频统计
        Map<String, Long> collect = data.stream()
                .collect(Collectors.groupingBy(x -> x, Collectors.counting()));

        // 根据词频生成权重
        for (String d : data) {
            wordFrequencies.add(new WordFrequency(d, Math.toIntExact(collect.get(d))));
        }
        //此处不设置会出现中文乱码
        java.awt.Font font = new java.awt.Font("STSong-Light", 2, 18);
        final Dimension dimension = new Dimension(600, 600);
        final WordCloud wordCloud = new WordCloud(dimension, CollisionMode.RECTANGLE);
        // 词间距
        wordCloud.setPadding(2);
        // 设置背景颜色
        wordCloud.setBackground(new CircleBackground(255));
        // 字体大小
        wordCloud.setFontScalar(new SqrtFontScalar(12, 42));
        //设置词云显示的三种颜色，越靠前设置表示词频越高的词语的颜色
        wordCloud.setColorPalette(color());
        // 生成字体
        wordCloud.setKumoFont(new KumoFont(font));
        wordCloud.setBackgroundColor(new Color(255, 255, 255));
        //因为我这边是生成一个圆形,这边设置圆的半径
        wordCloud.setBackground(new CircleBackground(255));
        wordCloud.build(wordFrequencies);
        ByteArrayOutputStream stream = new ByteArrayOutputStream();
        wordCloud.writeToStreamAsPNG(stream);

        return stream;
    }

    /**
     * 随机生成颜色
     * @return
     */
    private static ColorPalette color() {
        Random random = ThreadLocalRandom.current();
        Color[] colors = new Color[10];

        for (int i = 0; i < colors.length; ++i) {
            colors[i] = new Color(random.nextInt(230) + 25, random.nextInt(230) + 25, random.nextInt(230) + 25);
        }

        return new ColorPalette(colors);
    }

}
